#include "hla_processor.h"

class DispersionProcessor : public HlaGroupsProcessor
{
public:
    DispersionProcessor(){}
private:
    map<string, size_t> haplotype_rows;
    map<string, size_t> group_columns;
    map<size_t, map<size_t, float> > matrix;

    static size_t ensure_containment(map<string, size_t>& container, const string& element)
    {
        const map<string, size_t>::const_iterator it = container.find(element);

        if (it == container.end())
        {
            const size_t orig_size = container.size();
            container[element] = orig_size;
            return orig_size;
        }
        else
            return it->second;
    }

    void print_header(Sheet& sheet) const;

    void set_dispersion_data(const string& group, const string& haplotype, size_t unique_peptides, size_t total_peptides)
    {
        //cerr << group << "," << haplotype << ": " << unique_peptides << ", " << total_peptides << endl;
        matrix
            [ensure_containment(haplotype_rows, haplotype)]
            [ensure_containment(group_columns, group)]
        = total_peptides > 0 ? float(total_peptides - unique_peptides) / total_peptides : -1.0f;
    }

    virtual ProcessorPriority get_priority() const  {   return 2.0f;    }
    virtual void process(const Group& group);

    virtual void generate_results(Spreadsheet& spreadsheet);
};

REGISTER_PROCESSOR(DispersionProcessor);

void DispersionProcessor::process(const Group& group)
{
    size_t seq_count = 0;
    FirstTimeFlag first_haplotype;

    // for each haplotype in this group, analyze the peptides of each sequence:
    for (list<Haplotype>::const_iterator haplotype = group.haplotypes.begin(); haplotype != group.haplotypes.end(); ++haplotype)
    {
        set<string> peptides;
        size_t total_peptides = 0;

        // for each sequence in this haplotype, analyze its peptides
        for (list<Sequence>::const_iterator sequence = haplotype->sequences.begin(); sequence != haplotype->sequences.end(); ++sequence)
        {
            peptides.insert(sequence->peptides.begin(), sequence->peptides.end());
            total_peptides += sequence->peptides.size();
        }

        if (first_haplotype)
            seq_count = haplotype->sequences.size();
        else
        {
            if (seq_count != haplotype->sequences.size())
                cerr << "Error: different number of sequences among haplotypes in group " << group.name << endl;
        }

        set_dispersion_data(group.name, haplotype->name, peptides.size(), total_peptides);
    }
}

void DispersionProcessor::print_header(Sheet& sheet) const
{
    Row data(sheet);

    data << "";

    for (map<string, size_t>::const_iterator column = group_columns.begin(); column != group_columns.end(); ++column)
        data << column->first;
}

void DispersionProcessor::generate_results(Spreadsheet& spreadsheet)
{
    Sheet new_sheet( spreadsheet, "Dispersion" );

    print_header(new_sheet);

    for (map<string, size_t>::const_iterator row = haplotype_rows.begin(); row != haplotype_rows.end(); ++row)
    {
        Row data(new_sheet);

        data << row->first;

        for (map<string, size_t>::const_iterator column = group_columns.begin(); column != group_columns.end(); ++column)
            data << find(find(matrix, row->second), column->second);
    }
}

